rm(list=ls(all=TRUE))

library(MASS)
library(mc2d)

## prob. of 1 for randomizing device or group status indicator ##

p <- 0.25

n.MC <- 2000

pars.MC <- matrix(NA,nrow=n.MC,ncol=6)
mu.hat <- rep(NA,length.out=n.MC) 
delta  <- rep(NA,length.out=n.MC)


n.sim <- 5000  ## sample size for each Monte Carlo simulation ##

X.mu  <- c(.5, 1.5)
X.sig <- rbind(c(.1,.1),c(.1,.2))

lamT1.act <- 0.4
lamL1.act <- 0.4
lamT0.act <- 0.8
Beta.act  <- c(-1,2,-1.3)

for (k in 1:n.MC) {

X<-mvrnorm(n.sim,X.mu,X.sig)
X<-cbind(1,X)

mu.act <- (1+exp(-X%*%Beta.act))^-1
theta  <- rbinom(n=n.sim,size=1,mu.act)

gam1 <- p*lamT0.act*(1-theta) + (1-p)*lamL1.act*theta
gam2 <- (1-p)*lamT0.act*(1-theta) + p*lamL1.act*theta
gam3 <- (1-p)*lamT1.act*theta
gam4 <- p*lamT1.act*theta
gam5 <- p*(1-lamT0.act)*(1-theta) + (1-p)*(1-lamT1.act-lamL1.act)*theta
gam6 <- (1-p)*(1-lamT0.act)*(1-theta) + p*(1-lamT1.act-lamL1.act)*theta

Mprob <- cbind(gam1,gam2,gam3,gam4,gam5,gam6)

Y <- rmultinomial(n=n.sim,size=1,Mprob)
Y <- 1*Y[,1] + 2*Y[,2] + 3*Y[,3] + 4*Y[,4] + 5*Y[,5] + 6*Y[,6]

n.A <- sum(Y==3|Y==4)

## Set up function to calculate parameter values according to E-M algorithm ##


EM <- function(par) { 

lamT1 <-par[1]
lamL1 <-par[2]
lamT0 <-par[3]
Beta  <-par[4:length(par)]

mu    <- (1+exp(-X%*%Beta))^-1

Pr.1.3  <- as.numeric(Y>=3)*0 + as.numeric(Y==1)*((p*lamT0*(1-mu))/(p*lamT0*(1-mu) + (1-p)*lamL1*mu)) + as.numeric(Y==2)*(((1-p)*lamT0*(1-mu))/((1-p)*lamT0*(1-mu) + p*lamL1*mu))
Pr.2.4  <- as.numeric(Y>=3)*0 + as.numeric(Y<3)*(1-Pr.1.3)
Pr.5.6  <- ifelse(Y==3|Y==4,1,0)
Pr.7.9  <- as.numeric(Y<=4)*0 + as.numeric(Y==5)*((p*(1-lamT0)*(1-mu))/(p*(1-lamT0)*(1-mu) + (1-p)*(1-lamT1-lamL1)*mu)) + as.numeric(Y==6)*(((1-p)*(1-lamT0)*(1-mu))/((1-p)*(1-lamT0)*(1-mu) + p*(1-lamT1-lamL1)*mu))
Pr.8.10 <- as.numeric(Y<=4)*0 + as.numeric(Y>4)*(1-Pr.7.9)

Pr.T <- Pr.2.4 + Pr.5.6 + Pr.8.10

lamT1.j <- n.A/(n.A + sum(Pr.2.4) + sum(Pr.8.10) )
lamL1.j <- sum(Pr.2.4)/(n.A + sum(Pr.2.4) + sum(Pr.8.10) )
lamT0.j <- sum(Pr.1.3)/(sum(Pr.1.3)+sum(Pr.7.9) )
  
V <- mu*(1-mu)
X.til<-matrix(NA,nrow=nrow(X),ncol=ncol(X))
for (i in 1:n.sim) {X.til[i,]<-V[i]*X[i,]} 

D <- Pr.T - mu

Beta.j <- Beta + solve(t(X)%*%X.til,tol=.Machine$double.eps^2)%*%t(X)%*%D

est.j <- c(lamT1.j,lamL1.j,lamT0.j,Beta.j)
est.j}

tol <- .Machine$double.eps^.5 # tolerance for stopping algorithm

M<-4000
K<-ncol(X)
init<-c(.2,.2,.2,rep(0,K)) # initial parameter estimates
pars.old <- EM(init)


for (j in 1:M){

pars.new <- EM(pars.old)

abs.diff <- abs(pars.new-pars.old)

if (max(abs.diff) < tol ) break

pars.old <- pars.new}

pars.MC[k,] <- pars.new

X2.l <- quantile(X[,3],.05)
X2.h <- quantile(X[,3],.95)

X.1 <- cbind(X[,1:2],X2.h)
X.0 <- cbind(X[,1:2],X2.l)

B.new <- pars.new[4:6]

mu.ca   <- (1+exp(-X%*%B.new))^-1
mu.1    <- (1+exp(-X.1%*%B.new))^-1
mu.0    <- (1+exp(-X.0%*%B.new))^-1 

mu.hat[k]<-mean(mu.ca)
delta[k]<-mean(mu.1-mu.0)

}

lamT1.avg<-mean(pars.MC[,1])
lamL1.avg<-mean(pars.MC[,2])
lamT0.avg<-mean(pars.MC[,3])
beta0.avg<-mean(pars.MC[,4])
beta1.avg<-mean(pars.MC[,5])
beta2.avg<-mean(pars.MC[,6])

lamT1.bias <-lamT1.avg-lamT1.act
lamL1.bias <-lamL1.avg-lamL1.act
lamT0.bias <-lamT0.avg-lamT0.act
beta0.bias <-beta0.avg-Beta.act[1]
beta1.bias <-beta1.avg-Beta.act[2]
beta2.bias <-beta2.avg-Beta.act[3]

lamT1.mse <-mean((pars.MC[,1]-lamT1.act)^2)
lamL1.mse <-mean((pars.MC[,2]-lamL1.act)^2)
lamT0.mse <-mean((pars.MC[,3]-lamT0.act)^2)
beta0.mse <-mean((pars.MC[,4]-Beta.act[1])^2)
beta1.mse <-mean((pars.MC[,5]-Beta.act[2])^2)
beta2.mse <-mean((pars.MC[,6]-Beta.act[3])^2)

PARAMETERS<-c(lamT1.avg,lamL1.avg,lamT0.avg,beta0.avg,beta1.avg,beta2.avg)
BIAS<-c(lamT1.bias,lamL1.bias,lamT0.bias,beta0.bias,beta1.bias,beta2.bias)
MSE<-c(lamT1.mse,lamL1.mse,lamT0.mse,beta0.mse,beta1.mse,beta2.mse)

QOIs<-cbind(mu.hat,delta)

save(QOIs,file="JointQOIsHE5000.RData")


